package main

import "fmt"
import "runtime"
import "strconv"
import "flag"
import "log"
import "os"
import "time"
import "net"
import "net/http"
import "bufio" //用于逐行读取文本

// URL列表
var urlList chan string = make(chan string)

// 检测结果
var result chan string = make(chan string)

// 线程数
var thread chan int = make(chan int)
var workThread int

// 关闭程序
var clo chan bool = make(chan bool)

// 保存结果
func writeResult() {
	fileName := "result.txt"
	fout, err := os.Create(fileName)
	if err != nil {
		// 文件创建失败
		fmt.Println(fileName + " create failed!")
	}
	defer fout.Close()
	s, ok := <-result
	for ok {
		time.Sleep(50 * time.Millisecond)
		fout.WriteString(s + "\r\n")
		s, ok = <-result
	}
	// 等待结果保存完成后通知进程退出
	clo <- true
}

// 根据线程参数启动检测器
func runChecker() {
	t, ok := <-thread
	workThread = t
	if ok {
		for i := 1; i <= workThread; i++ {
			go checker(strconv.Itoa(i))
		}
	}
	// 等待线程终止
	for <-thread == -1 {
		workThread--
		if workThread == 0 {
			// 全部线程已终止，关闭结果写入，退出程序
			close(result)
			break
		}
	}
}

func dialTimeout(network, addr string) (net.Conn, error) {
	var timeout = time.Duration(20000 * time.Millisecond)
	return net.DialTimeout(network, addr, timeout)
}

// 检测线程
func checker(threadId string) {
	url, ok := <-urlList
	fmt.Println("[检测线程-" + threadId + "] 启动")
	for ok {
		// 校验url
		if url == "" {
			log.Printf("[空白]")
		} else {
			transport := http.Transport{
				Dial:              dialTimeout,
				DisableKeepAlives: true,
			}
			client := http.Client{
				Transport: &transport,
				Timeout:   time.Duration(15 * time.Second),
			}
			response, err := client.Head(url)
			if err != nil {
				log.Printf("[检测出错]：\r\n%q\r\n线程：%q", err, threadId)
				//fmt.Println("NumGoroutine ", runtime.NumGoroutine());
			} else {
				if response.StatusCode == 200 {
					// URL可访问
					result <- url
				}
				response.Body.Close()
			}
		}
		url, ok = <-urlList
	}
	fmt.Println("[检测线程-" + threadId + "] 完成")
	thread <- -1
}

// 读取URL列表文件
func getUrls(fileName string) {
	fin, err := os.Open(fileName)
	defer fin.Close()
	if err != nil {
		fmt.Println(err)
		return
	}
	buff := bufio.NewReader(fin)
	for {
		// 逐行读取，没有多余字符
		line, _, err := buff.ReadLine()
		// 信道超时设置
		t := time.NewTimer(150 * time.Minute)
		defer t.Stop()
		select {
		case urlList <- string(line):
		case <-t.C:
			fmt.Println("写入URL地址信道超时")
		}
		if err != nil {
			break
		}
	}
	close(urlList)
}

// 处理参数
func processFlag(arg []string) {
	// 线程数
	t, err := strconv.Atoi(arg[1])
	if err != nil {
		// 线程数设置错误
		fmt.Println("'thread' Setting error")
		return
	}
	// 限制最大线程数4096
	if t < 1 {
		t = 1
	} else if t > 4096 {
		t = 4096
	}
	// 传递启动进程数
	thread <- t

	// URLs列表文件
	fileName := arg[0]
	// 传递URL列表
	getUrls(fileName)

}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()
	if flag.NArg() != 2 {
		// 参数不合法
		fmt.Println("Usage: checkURL URL.lst threadNumber")
		fmt.Println("Author: Frank Kung <i@kongfanjian.com>")
		return
	}

	// 获取参数
	args := make([]string, 0, 2)
	for i := 0; i < flag.NArg(); i++ {
		args = append(args, flag.Arg(i))
	}

	// 启动检测线程
	go runChecker()
	// 启动结果写入线程
	go writeResult()
	// 参数处理
	processFlag(args)

	// 等待退出指令
	<-clo
	fmt.Println("Exit")
}
